package com.etl.dataflow.scheduler.admin.service;

import cn.hutool.core.bean.BeanUtil;
import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.http.HttpResponse;
import cn.hutool.http.HttpUtil;
import com.baomidou.mybatisplus.core.conditions.query.QueryWrapper;
import com.baomidou.mybatisplus.extension.plugins.pagination.Page;
import com.baomidou.mybatisplus.extension.service.impl.ServiceImpl;
import com.etl.dataflow.common.entity.ExecutorInfo;
import com.etl.dataflow.common.entity.TaskInfo;
import com.etl.dataflow.common.enums.TaskStateEnum;
import com.etl.dataflow.common.exception.EtlException;
import com.etl.dataflow.common.exception.ServiceException;
import com.etl.dataflow.common.response.Code;
import com.etl.dataflow.common.response.R;
import com.etl.dataflow.common.util.ResponseUtil;
import com.etl.dataflow.scheduler.admin.entity.DataflowDatasource;
import com.etl.dataflow.scheduler.admin.entity.DataflowTaskGroupInfo;
import com.etl.dataflow.scheduler.admin.entity.DataflowTaskInfo;
import com.etl.dataflow.scheduler.admin.entity.DataflowTaskRecord;
import com.etl.dataflow.scheduler.admin.entity.request.PageRequest;
import com.etl.dataflow.scheduler.admin.entity.vo.DataflowTaskInfoVo;
import com.etl.dataflow.scheduler.admin.mapper.DataflowDatasourceMapper;
import com.etl.dataflow.scheduler.admin.mapper.DataflowTaskGroupInfoMapper;
import com.etl.dataflow.scheduler.admin.mapper.DataflowTaskInfoMapper;
import com.etl.dataflow.scheduler.admin.mapper.DataflowTaskRecordMapper;
import com.etl.dataflow.scheduler.admin.schedule.TaskConcurrencyManager;
import com.etl.dataflow.scheduler.communication.ExecutorStateChecker;
import com.etl.dataflow.scheduler.cron.CronExpression;
import com.etl.dataflow.scheduler.timewheel.TimeWheelFacade;
import com.etl.dataflow.scheduler.timewheel.WheelTask;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;

import javax.annotation.PostConstruct;
import java.text.ParseException;
import java.util.*;
import java.util.function.Function;
import java.util.stream.Collectors;

/**
 * <p>
 * 任务信息 服务实现类
 * </p>
 *
 * @author dx
 * @since 2022-04-24
 */
@Slf4j
@Service
public class DataflowTaskInfoService extends ServiceImpl<DataflowTaskInfoMapper, DataflowTaskInfo>
        implements DefaultService<DataflowTaskInfo> {

    private final DataflowTaskRecordMapper taskRecordMapper;
    private final DataflowDatasourceMapper datasourceMapper;
    private final DataflowTaskInfoMapper taskInfoMapper;
    private final DataflowTaskGroupInfoMapper taskGroupMapper;
    private final ExecutorStateChecker executorStateChecker;
    private final TaskConcurrencyManager taskConcurrencyManager;

    public DataflowTaskInfoService(DataflowTaskRecordMapper taskRecordMapper,
                                   DataflowDatasourceMapper datasourceMapper, DataflowTaskInfoMapper taskInfoMapper, DataflowTaskGroupInfoMapper taskGroupMapper, ExecutorStateChecker executorStateChecker, TaskConcurrencyManager taskConcurrencyManager) {
        this.taskRecordMapper = taskRecordMapper;
        this.datasourceMapper = datasourceMapper;
        this.taskInfoMapper = taskInfoMapper;
        this.taskGroupMapper = taskGroupMapper;
        this.executorStateChecker = executorStateChecker;
        this.taskConcurrencyManager = taskConcurrencyManager;
    }

    @PostConstruct
    private void loadEnabledTimerTaskIntoTimeWheel() {
        QueryWrapper<DataflowTaskInfo> queryWrapper = new QueryWrapper<>();
        queryWrapper.eq("timer_state", DataflowTaskInfo.TimerState.ENABLED.ordinal());
        List<DataflowTaskInfo> enabledTasks = taskInfoMapper.selectList(queryWrapper);
        TimeWheelFacade timeWheelFacade = TimeWheelFacade.getInstance();
        for (DataflowTaskInfo enabledTask : enabledTasks) {
            // 从时间轮删除旧任务
            timeWheelFacade.removeTask(String.valueOf(enabledTask.getId()));
            // 重新加入时间轮
            WheelTask wheelTask = getWheelTask(enabledTask.getId(), enabledTask.getName());
            try {
                timeWheelFacade.addTask(new CronExpression(enabledTask.getCronExpr()), wheelTask);
            } catch (ParseException ignored) {
            }
        }
    }

    public R addOrUpdate(DataflowTaskInfo taskInfo) {
        CronExpression newCronExpression = parseCronExpression(taskInfo);

        if (null == taskInfo.getId()) {
            taskInfo.setTimerState(DataflowTaskInfo.TimerState.DISABLED.ordinal());
            super.save(taskInfo);
            return R.success();
        }

        DataflowTaskInfo oldTaskInfo = taskInfoMapper.selectById(taskInfo.getId());
        if (oldTaskInfo == null) {
            super.save(taskInfo);
            return R.success();
        }

        boolean isTimerEnabled = oldTaskInfo.getTimerState() == DataflowTaskInfo.TimerState.ENABLED.ordinal();
        boolean isCronExprUpdated = !StrUtil.equals(oldTaskInfo.getCronExpr(), taskInfo.getCronExpr());
        if (isTimerEnabled && isCronExprUpdated) {
            // 从时间轮删除旧任务
            TimeWheelFacade.getInstance().removeTask(String.valueOf(taskInfo.getId()));
            // 重新加入时间轮
            WheelTask wheelTask = getWheelTask(taskInfo.getId(), taskInfo.getName());
            TimeWheelFacade.getInstance().addTask(newCronExpression, wheelTask);
        }
        super.updateById(taskInfo);
        return R.success();
    }

    private WheelTask getWheelTask(Long taskId, String taskName) {
        return new WheelTask(String.valueOf(taskId)) {
            @Override
            public void run() {
                log.debug(">>>>>>>>>>>>触发任务: " + taskName + ", taskId: " + taskId);
                runTask(taskId);
            }
        };
    }

    private CronExpression parseCronExpression(DataflowTaskInfo taskInfo) {
        CronExpression newCronExpression;
        try {
            newCronExpression = new CronExpression(taskInfo.getCronExpr());
        } catch (ParseException e) {
            throw new ServiceException(Code.CRON_EXPR_ERROR);
        }
        return newCronExpression;
    }

    public Page<? extends DataflowTaskInfo> page(PageRequest pageRequest) {
        Page<DataflowTaskInfo> page = super.baseMapper.selectPage(toPage(pageRequest), orLikeAllFields(pageRequest, DataflowTaskInfo.class, "id"));
        if (page.getTotal() <= 0 || page.getRecords().size() <= 0) {
            return page;
        }
        Map<Long, DataflowDatasource> srcDsMap = queryDatasource(page, DataflowTaskInfo::getSrcDsId);
        Map<Long, DataflowDatasource> dstDsMap = queryDatasource(page, DataflowTaskInfo::getDstDsId);
        Map<Long, DataflowTaskGroupInfo> taskGroupInfoMap = queryTaskGroup(page, DataflowTaskInfo::getGroupId);
        List<DataflowTaskInfoVo> vos = new ArrayList<>();
        for (DataflowTaskInfo taskInfo : page.getRecords()) {
            DataflowTaskInfoVo vo = new DataflowTaskInfoVo();
            BeanUtil.copyProperties(taskInfo, vo);
            vo.setSrcDatasource(srcDsMap.get(taskInfo.getSrcDsId()));
            vo.setDstDatasource(dstDsMap.get(taskInfo.getDstDsId()));
            vo.setTaskGroupInfo(taskGroupInfoMap.get(taskInfo.getGroupId()));
            vos.add(vo);
        }
        Page<DataflowTaskInfoVo> pageVo = new Page<>();
        BeanUtil.copyProperties(page, pageVo);
        pageVo.setRecords(vos);
        return pageVo;
    }

    private Map<Long, DataflowTaskGroupInfo> queryTaskGroup(Page<DataflowTaskInfo> page, Function<DataflowTaskInfo, Long> getGroupId) {
        List<Long> taskGroupIds = page.getRecords().stream().map(getGroupId).collect(Collectors.toList());
        QueryWrapper<DataflowTaskGroupInfo> queryWrapper = new QueryWrapper<>();
        return taskGroupMapper.selectList(queryWrapper.in("id", taskGroupIds)).stream().collect(Collectors.toMap(DataflowTaskGroupInfo::getId, v -> v));
    }

    private Map<Long, DataflowDatasource> queryDatasource(Page<DataflowTaskInfo> page, Function<DataflowTaskInfo, Long> getDsId) {
        List<Long> dsIds = page.getRecords().stream().map(getDsId).collect(Collectors.toList());
        QueryWrapper<DataflowDatasource> queryWrapper = new QueryWrapper<>();
        return datasourceMapper.selectList(queryWrapper.in("id", dsIds)).stream().collect(Collectors.toMap(DataflowDatasource::getId, v -> v));
    }

    public R runTask(Long taskId) {
        // 先检查任务是否在等待队列中
        if (taskConcurrencyManager.isTaskWaiting(taskId)) {
            throw new ServiceException(Code.TASK_EXISTS);
        }
        // 检查任务是否正在运行中
        this.checkIsTaskExists(taskId + "");

        DataflowTaskInfo taskInfo = taskInfoMapper.selectById(taskId);
        // 将任务放入等待队列
        taskConcurrencyManager.submitTask(taskInfo);
        return R.success();
    }

    private void checkIsTaskExists(String taskId) {
        for (ExecutorInfo onlineExecutor : executorStateChecker.getOnlineExecutors()) {
            for (TaskInfo task : onlineExecutor.getTaskList()) {
                if (StrUtil.equals(task.getId(), taskId)) {
                    throw new ServiceException(Code.TASK_EXISTS);
                }
            }
        }
    }

    /**
     * 执行全部任务
     *
     * @return R
     */
    public R runAllTask() {
        List<DataflowTaskInfo> taskInfos = taskInfoMapper.selectList(null);

        for (DataflowTaskInfo taskInfo : taskInfos) {
            try {
                taskConcurrencyManager.submitTask(taskInfo);
            } catch (EtlException e) {
                log.error(e.getMessage());
            }
        }
        return R.success();
    }

    public R runSpecificTasks(List<Long> taskIds) {
        List<DataflowTaskInfo> taskInfos = taskInfoMapper.selectBatchIds(taskIds);
        String msg = null;
        for (DataflowTaskInfo taskInfo : taskInfos) {
            try {
                taskConcurrencyManager.submitTask(taskInfo);
            } catch (EtlException e) {
                log.error(e.getMessage());
                if (Optional.ofNullable(e.getMessage()).orElse("").startsWith("任务重复提交")) {
                    msg = "运行成功，存在重复任务，已过滤";
                }
            }
        }
        return msg == null ? R.success() : R.success().setMsg(msg);
    }

    @Transactional(rollbackFor = Exception.class)
    public R enableTimer(Long taskId) throws ParseException {
        DataflowTaskInfo taskInfo = this.updateTimerState(taskId, DataflowTaskInfo.TimerState.ENABLED.ordinal());
        WheelTask wheelTask = getWheelTask(taskId, taskInfo.getName());
        TimeWheelFacade.getInstance().addTask(new CronExpression(taskInfo.getCronExpr()), wheelTask);
        return R.success();
    }

    @Transactional(rollbackFor = Exception.class)
    public R enableSpecificTimers(List<Long> taskIds) throws ParseException {
        int effectRows = taskInfoMapper.updateTimerStateByIds(taskIds, DataflowTaskInfo.TimerState.ENABLED.ordinal());
        if (effectRows <= 0) {
            throw new ServiceException(Code.CANNOT_ENABLE_TIMER);
        }
        Map<Long, DataflowTaskInfo> taskMap = taskInfoMapper.selectBatchIds(taskIds).stream()
                .collect(Collectors.toMap(DataflowTaskInfo::getId, v -> v, (t, t2) -> t2));
        for (Long taskId : taskIds) {
            WheelTask wheelTask = getWheelTask(taskId, taskMap.get(taskId).getName());
            TimeWheelFacade.getInstance().addTask(new CronExpression(taskMap.get(taskId).getCronExpr()), wheelTask);
        }
        return R.success();
    }

    @Transactional(rollbackFor = Exception.class)
    public R disableSpecificTimers(List<Long> taskIds) {
        int effectRows = taskInfoMapper.updateTimerStateByIds(taskIds, DataflowTaskInfo.TimerState.DISABLED.ordinal());
        if (effectRows <= 0) {
            throw new ServiceException(Code.CANNOT_DISABLE_TIMER);
        }
        for (Long taskId : taskIds) {
            TimeWheelFacade.getInstance().removeTask(String.valueOf(taskId));
        }
        return R.success();
    }

    @Transactional(rollbackFor = Exception.class)
    public R disableTimer(Long taskId) {
        this.updateTimerState(taskId, DataflowTaskInfo.TimerState.DISABLED.ordinal());
        TimeWheelFacade.getInstance().removeTask(String.valueOf(taskId));
        return R.success();
    }

    private DataflowTaskInfo updateTimerState(Long taskId, int state) {
        DataflowTaskInfo taskInfo = super.getById(taskId);
        taskInfo.setTimerState(state);
        super.updateById(taskInfo);
        return taskInfo;
    }

    @Transactional(rollbackFor = Exception.class)
    public R disableAllTimer() {
        QueryWrapper<DataflowTaskInfo> queryWrapper = new QueryWrapper<>();
        queryWrapper.eq("timer_state", "1");
        List<DataflowTaskInfo> enabledTasks = super.baseMapper.selectList(queryWrapper);
        if (CollUtil.isEmpty(enabledTasks)) {
            throw new ServiceException(Code.NO_TASK_NEED_DISABLE_TIMER);
        }
        TimeWheelFacade timeWheelFacade = TimeWheelFacade.getInstance();
        for (DataflowTaskInfo enabledTask : enabledTasks) {
            enabledTask.setTimerState(DataflowTaskInfo.TimerState.DISABLED.ordinal());
        }
        boolean success = super.updateBatchById(enabledTasks);
        if (success) {
            for (DataflowTaskInfo enabledTask : enabledTasks) {
                // 从时间轮移除任务
                timeWheelFacade.removeTask(String.valueOf(enabledTask.getId()));
            }
            return R.success();
        }
        throw new ServiceException(Code.CANNOT_DISABLE_TIMER);
    }

    public R terminateTask(Long taskRecordId, String uuid, Long taskId) {
        DataflowTaskRecord taskRecord = taskRecordMapper.selectById(taskRecordId);
        String taskIdStr = taskId + "";
        ExecutorInfo executorInfo = executorStateChecker.getExecutorInfoByUuid(uuid);
        Iterator<TaskInfo> iterator = executorInfo.getTaskList().iterator();
        while (iterator.hasNext()) {
            TaskInfo taskInfo = iterator.next();
            if (StrUtil.equals(taskInfo.getId(), taskIdStr)) {
                String body;
                try (HttpResponse httpResponse = HttpUtil
                        .createGet(executorInfo.getExecutorUrl() + "/task/terminate/" + taskId)
                        .execute()) {
                    body = httpResponse.body();
                }
                log.debug(body);
                if (ResponseUtil.isSuccess(body) || ResponseUtil.codeEquals(body, Code.TASK_NOT_FOUND_CANNOT_TERMINATE)) {
                    // 从内存中删除此任务
                    iterator.remove();
                    taskRecord.setTaskState(TaskStateEnum.TERMINATED.getValue());
                    taskRecord.setEndTime(new Date());
                    taskRecordMapper.updateById(taskRecord);
                    return R.success();
                }
                return new R(Code.TASK_TERMINATE_ERROR).put("data", body);
            }
        }
        taskRecord.setTaskState(TaskStateEnum.FATAL_ERROR.getValue());
        taskRecordMapper.updateById(taskRecord);
        return new R(Code.TASK_NOT_FOUND_CANNOT_TERMINATE);
    }

    public R getChildren(Long groupId) {
        QueryWrapper<DataflowTaskInfo> queryWrapper = new QueryWrapper<>();
        queryWrapper.eq("group_id", groupId);
        List<DataflowTaskInfo> taskInfos = super.baseMapper.selectList(queryWrapper);
        return R.success().put("data", taskInfos);
    }

}
